import subprocess
import re
import os
from collections import defaultdict
from tools.logger import logger
from tools.generator.classifier import (
    SLICE_TYPE_BINS,
    SLICE_TYPE_LIBS,
    SLICE_TYPE_CONFIG,
)


class DependencyAnalyzer:
    """
    Analyzes binary dependencies for a given package.
    """

    def __init__(
        self,
        package_name: str,
        pkg_extract_path: str,
        classified_slices: dict[str, set[str]],
    ):
        self.package_name = package_name
        self.pkg_extract_path = pkg_extract_path
        self.classified_slices = classified_slices

        # Caches to improve performance
        self._file_to_pkg_cache = {}
        self._ldconfig_cache = None

    def analyze(self) -> dict[str, set[str]]:
        """
        Main entry point for the analysis.
        Orchestrates the process of analyzing dependencies for all relevant slices.

        Returns:
        - A dictionary mapping slice names to sets of dependent slices.
        """
        # Step 1: Analyze external binary dependencies (existing logic)
        slice_deps = self._analyze_external_dependencies()

        # Step 2: Inject internal dependencies (new logic)
        self._inject_internal_dependencies(slice_deps)

        return slice_deps

    def _analyze_external_dependencies(self) -> dict[str, set[str]]:
        """
        Analyzes dependencies on external packages using readelf.
        This method checks for shared libraries used by binaries
        and maps them to their owning packages.
        Returns:
        - A dictionary mapping slice names to sets of dependent slices.
        """

        slice_deps = defaultdict(set)

        # Pre-load the ldconfig cache once for the entire analysis
        self._load_ldconfig_cache()

        # Iterate through classified slices to find binaries and libraries
        for slice_name, file_set in self.classified_slices.items():
            if not (slice_name.endswith(SLICE_TYPE_BINS) or slice_name.endswith(SLICE_TYPE_LIBS)):
                continue

            for file in file_set:
                full_file_path = os.path.join(self.pkg_extract_path, file.lstrip("/"))

                if not os.path.isfile(full_file_path):
                    continue

                # Get dependencies for a single binary file
                needed_libs = self._get_needed_libraries(full_file_path)

                for lib_name in needed_libs:
                    # Resolve the library name to a package owner
                    owner_pkg = self._resolve_library_to_package(lib_name)

                    if owner_pkg and owner_pkg != self.package_name:
                        dep_slice = f"{owner_pkg}_libs"
                        logger.debug(
                            f"'{slice_name}' dependency found: {file} -> {lib_name} ({dep_slice})"
                        )
                        slice_deps[slice_name].add(dep_slice)

        return slice_deps

    def _inject_internal_dependencies(self, slice_deps: dict[str, set[str]]):
        """
        Adds dependencies between slices of the same package (e.g., bins -> config).
        """
        logger.debug("Injecting internal dependencies...")

        # Construct the names of the potential slices within this package
        config_slice_name = f"{self.package_name}{SLICE_TYPE_CONFIG}"
        bins_slice_name = f"{self.package_name}{SLICE_TYPE_BINS}"
        libs_slice_name = f"{self.package_name}{SLICE_TYPE_LIBS}"

        # Check if a config slice exists for this package
        if config_slice_name in self.classified_slices:
            # If bins slice exists, make it depend on the config slice
            if bins_slice_name in self.classified_slices:
                logger.info(f"Adding internal dependency: {bins_slice_name} -> {config_slice_name}")
                slice_deps[bins_slice_name].add(config_slice_name)

            # If libs slice exists, make it depend on the config slice
            if libs_slice_name in self.classified_slices:
                logger.info(f"Adding internal dependency: {libs_slice_name} -> {config_slice_name}")
                slice_deps[libs_slice_name].add(config_slice_name)

    def _load_ldconfig_cache(self):
        """
        Executes `ldconfig -p` once and caches its output.
        """
        logger.debug("Loading ldconfig cache...")
        try:
            ldconfig_proc = subprocess.run(
                ["ldconfig", "-p"], capture_output=True, text=True, check=True
            )
            self._ldconfig_cache = ldconfig_proc.stdout
        except FileNotFoundError as e:
            raise RuntimeError(
                "`ldconfig` command not found. Please ensure the `glibc` package is installed."
            ) from e
        except subprocess.CalledProcessError as e:
            raise RuntimeError(
                f"`ldconfig` command failed: {e.stderr.strip()}. Halting analysis."
            ) from e

    def _get_pkg_owner(self, file_path: str) -> str:
        """Finds the package that owns a file using an instance cache."""
        if file_path in self._file_to_pkg_cache:
            return self._file_to_pkg_cache[file_path]
        try:
            rpm_qf_proc = subprocess.run(
                ["rpm", "-qf", file_path], capture_output=True, text=True, check=True
            )
            pkg_full_name = rpm_qf_proc.stdout.strip()
            rpm_q_qf_proc = subprocess.run(
                ["rpm", "-q", "--qf", "%{NAME}", pkg_full_name],
                capture_output=True,
                text=True,
                check=True,
            )
            owner_pkg = rpm_q_qf_proc.stdout.strip()
            self._file_to_pkg_cache[file_path] = owner_pkg
            return owner_pkg
        except FileNotFoundError as e:
            raise RuntimeError(
                "FATAL: `rpm` command not found. Please ensure the `rpm` package is installed."
            ) from e
        except subprocess.CalledProcessError:
            self._file_to_pkg_cache[file_path] = ""
            return ""

    def _get_needed_libraries(self, binary_path: str) -> list[str]:
        """
        Runs `readelf -d` on a binary. Assumes the file is a valid ELF
        as it has been pre-filtered by the classifier.

        Returns:
            A list of needed libraries.
        """
        needed = []
        try:
            readelf_result = subprocess.run(
                ["readelf", "-d", binary_path],
                capture_output=True,
                text=True,
                check=True,
                encoding="utf-8",
                errors="ignore",
            )
            for line in readelf_result.stdout.strip().split("\n"):
                if "(NEEDED)" in line:
                    match = re.search(r"\[(.*)\]", line)
                    if match:
                        needed.append(match.group(1))
            return needed
        except FileNotFoundError as e:
            raise RuntimeError(
                "FATAL: `readelf` command not found. Please ensure the `binutils` "
                "package is installed."
            ) from e
        except subprocess.CalledProcessError as e:
            raise RuntimeError(
                f"readelf failed on '{binary_path}' "
                f"which was identified as an ELF file. The file might be corrupted. "
                f"Halting analysis. Error: {e.stderr.strip()}"
            ) from e

    def _resolve_library_to_package(self, lib_name: str) -> str:
        """
        Resolves a library name (e.g., 'libc.so.6') to its owning package name (e.g., 'glibc').
        Uses the ldconfig cache to find the library path.
        """
        if not self._ldconfig_cache:
            return ""

        lib_path_match = re.search(
            rf"\s+{re.escape(lib_name)}\s*.*=>\s*(/.*)", self._ldconfig_cache
        )
        if not lib_path_match:
            logger.warning(f"Could not find path for library '{lib_name}' in ldconfig cache.")
            return ""

        lib_path = lib_path_match.group(1)
        return self._get_pkg_owner(lib_path)
